#!/usr/bin/env python

import json
from src.data import pathparser
from tqdm.auto import tqdm


def splits_final_json_to_ymls(splits_final, ymls):
    with open(splits_final, 'r') as f:
        splits_final = json.load(f)
    for j, yml in tqdm(enumerate(ymls), total=len(ymls)):
        pathparser.saveyaml({
            'training': [{'image': ['image', i], 'label': ['label', i]} for i in splits_final[j]['train']],
            'validation': [{'image': ['image', i], 'label': ['label', i]} for i in splits_final[j]['val']]
        }, yml)


if __name__ == '__main__':
    splits_final_json_to_ymls(
        '/home/yusongli/Templates/yunet/nnUNet_preprocessed/Dataset002_C_intensity1500_roi2.0/splits_final.json',
        [
            '/home/yusongli/Templates/yunet/nnUNet_preprocessed/Dataset002_C_intensity1500_roi2.0/splits_final_f0.yml',
            '/home/yusongli/Templates/yunet/nnUNet_preprocessed/Dataset002_C_intensity1500_roi2.0/splits_final_f1.yml',
            '/home/yusongli/Templates/yunet/nnUNet_preprocessed/Dataset002_C_intensity1500_roi2.0/splits_final_f2.yml',
            '/home/yusongli/Templates/yunet/nnUNet_preprocessed/Dataset002_C_intensity1500_roi2.0/splits_final_f3.yml',
            '/home/yusongli/Templates/yunet/nnUNet_preprocessed/Dataset002_C_intensity1500_roi2.0/splits_final_f4.yml',
        ]
    )
